"""
mqtt 控制报文结构的处理

报文处理需要完成控制报文的创建和解析, 包括服务器和客户端所有类型的报文。
    控制报文分三个结构:
        Fixed header       固定报头, 所有控制报文都包含
        Variable header    可变报头, 部分控制报文包含
        Payload            有效载荷, 部分控制报文包含

    每个部分都需要处理报文的创建和解析即报文的转换，从各种类型转成字节或从
    字节转成相应格式的对象。正常情况下IO中保存的数据都是符合MQTT协议的数据，
    因此可以直接对缓冲区进行数据的读取和写入。因此三种结构对应三个类，而每个
    类都需要读写缓冲区，分别为 to_stream 和 from_stream 方法。

    FixedHeader 所有报文包含，并且包含三个信息，分别为: 包类型，类型的标志
    位，剩余长度。

    VariableHader 可变报头一般包含两个字节的报文标识符。

    Payload 有效载荷，对发布报文来说就是应用消息。

"""
from datetime import datetime
from struct import unpack

from umq.errors import MQTTException, CodecException, NoDataException
from umq.adapters import StreamReaderAdapter, StreamWriterAdapter
from umq.codecs import read_or_raise, int_to_bytes, decode_packet_id


# 控制报文的类型
# 保留
RESERVED_0 = 0x00
# 客户端到服务端, 客户端请求连接服务器
CONNECT = 0x01
# 服务端到客户端, 连接报文确认
CONNACK = 0x02
# 两端都允许, 发布消息
PUBLISH = 0x03
# 两端都允许, QoS 1 消息发布收到确认
PUBACK = 0x04
# 两端都允许, 发布收到(保证交付第一步)
PUBREC = 0x05
# 两端都允许, 发布释放(保证交付第二步)
PUBREL = 0x06
# 两端都允许, QoS 2 消息发布完成(保证交付第三步)
PUBCOMP = 0x07
# 客户端到服务端, 客户端订阅请求
SUBSCRIBE = 0x08
# 服务端到客户端, 订阅请求报文确认
SUBACK = 0x09
# 客户端到服务端, 客户端取消订阅请求
UNSUBSCRIBE = 0x0A
# 服务端到客户端, 取消订阅报文确认
UNSUBACK = 0x0B
# 客户端到服务端, 心跳请求
PINGREQ = 0x0C
# 服务端到客户端, 心跳响应
PINGRESP = 0x0D
# 客户端到服务端, 客户端断开连接
DISCONNECT = 0x0E
# 保留
RESERVED_15 = 0x0F


class MqttFixedHeader:
    """mqtt fixed header"""
    __slots__ = ('packet_type', 'flags', 'remaining_length')

    def __init__(self, packet_type, flag=0, length=0):
        """
        :packet_type: mqtt 控制报文类型
        :flag: 用于指定控制报文类型的标志位
        :length: 剩余长度
        """
        self.flags = flag
        self.packet_type = packet_type
        self.remaining_length = length

    def to_bytes(self):
        """将对象转成bytes"""
        def encode_remaining_length(length: int):
            """剩余长度是一个可变的，最大四个字节"""
            encoded = bytearray()
            while True:
                length_byte = length % 0x80
                length //= 0x80
                if length > 0:
                    length_byte |= 0x80
                encoded.append(length_byte)
                if length <= 0:
                    break
            return encoded

        out = bytearray()
        packet_type = 0
        try:
            packet_type = (self.packet_type << 4) | self.flags
            out.append(packet_type)
        except OverflowError:
            # 流溢出异常
            raise CodecException('packet type too long:value=%d'%packet_type)

        encoded_length = encode_remaining_length(self.remaining_length)
        out.extend(encoded_length)
        return out

    async def to_stream(self, writer: StreamWriterAdapter):
        """将报文写入缓冲区"""
        # 将对象转成字节流
        writer.write(self.to_bytes())

    @classmethod
    async def from_stream(cls, reader: StreamReaderAdapter):
        """
        从缓冲区获取数据
        :return: MqttFixedHeader实例
        """
        async def decode_remaining_length():
            """按照 MQTT 协议解码"""
            mul = 1
            value = 0
            read_buf = bytearray()
            while True:
                encoded_byte = await reader.read(1)
                int_byte = unpack('!B', encoded_byte)
                read_buf.append(int_byte[0])
                value += (int_byte[0] & 0x7F) * mul
                if (int_byte[0] & 0x80) == 0:
                    break
                else:
                    mul *= 0x80
                    if mul > 0x80 * 0x80 * 0x80:
                        raise MQTTException('Remaining too long')
            return value

        try:
            buf = await read_or_raise(reader, 1)
            int1 = unpack('!B', buf)
            msg_type = (int1[0] & 0xF0) >> 4
            flag = int1[0] & 0x0F
            remain_length = await decode_remaining_length()
            return cls(msg_type, flag, remain_length)
        except NoDataException:
            return None

    def __repr__(self):
        return type(self).__name__ + 'length={0}, flags={1}'.format(
            self.remaining_length, hex(self.flags)
        )


class MqttVariableHeader:
    """mqtt variable header"""

    def __init__(self):
        pass

    def to_bytes(self) -> bytes:
        """将对象转成bytes"""

    async def to_stream(self, writer: StreamWriterAdapter):
        """将报文写入缓冲区"""
        writer.write(self.to_bytes())
        await writer.drain()

    @classmethod
    async def from_stream(cls, reader: StreamReaderAdapter, fixed: MqttFixedHeader):
        """从缓冲区获取数据"""
        pass

class PackeIdVariableHeader(MqttVariableHeader):
    """包含packet id的可变报文头"""
    __slots__ = ('packet_id', )

    def __init__(self, packet_id):
        super().__init__()
        self.packet_id = packet_id

    def to_bytes(self):
        out = b''
        out += int_to_bytes(self.packet_id, 2)
        return out

    @classmethod
    async def from_stream(cls, reader: StreamReaderAdapter, fixed: MqttFixedHeader):
        packet_id = await decode_packet_id(reader)
        return cls(packet_id)

    def __repr__(self):
        return type(self).__name__ + '(packet_id+{})'.format(self.packet_id)


class MqttPayload:
    """mqtt payload"""

    def __init__(self):
        pass

    def to_bytes(self) -> bytes:
        """将对象转成bytes"""

    async def to_stream(self, writer: StreamWriterAdapter):
        """将报文写入缓冲区"""
        writer.write(self.to_bytes())
        await writer.drain()

    @classmethod
    async def from_stream(cls, reader: StreamReaderAdapter,
                          fixed: MqttFixedHeader,
                          variable: MqttVariableHeader):
        """从缓冲区获取数据"""
        pass

class MqttPacket:
    """mqtt packet"""
    __slots__ = ('fixed_header', 'variable_header', 'payload', 'protocol_ts')

    FIXED_HEADER = MqttFixedHeader
    VARIABLE_HEADER = None
    PAYLOAD = None

    def __init__(self, fixed: MqttFixedHeader,
                 variable: MqttVariableHeader = None,
                 payload: MqttPayload = None):
        self.fixed_header = fixed
        self.variable_header = variable
        self.payload = payload
        self.protocol_ts = None

    def to_bytes(self) -> bytes:
        """将对象转成bytes"""
        variable_header_bytes = b''
        payload_bytes = b''
        if self.variable_header:
            variable_header_bytes = self.variable_header.to_bytes()
        if self.payload:
            payload_bytes = self.payload.to_bytes()
        self.fixed_header.remaining_length = len(variable_header_bytes) + len(payload_bytes)
        fixed_header_bytes = self.fixed_header.to_bytes()
        return fixed_header_bytes + variable_header_bytes + payload_bytes

    async def to_stream(self, writer: StreamWriterAdapter):
        """将报文写入缓冲区"""
        writer.write(self.to_bytes())
        await writer.drain()
        self.protocol_ts = datetime.now()

    @classmethod
    async def from_stream(cls, reader: StreamReaderAdapter,
                          fixed_header=None, variable_header=None):
        """从缓冲区获取数据"""
        if fixed_header is None:
            fixed_header = await MqttFixedHeader.from_stream(reader)
        if cls.VARIABLE_HEADER:
            if variable_header is None:
                variable_header = await cls.VARIABLE_HEADER.from_stream(
                    reader, fixed_header
                )
        else:
            variable_header = None
        if cls.PAYLOAD:
            payload = await cls.PAYLOAD.from_stream(
                reader, fixed_header, variable_header
            )
        else:
            payload = None

        if fixed_header and not variable_header and not payload:
            instance = cls(fixed_header)
        elif fixed_header and not payload:
            instance = cls(fixed_header, variable_header)
        else:
            instance = cls(fixed_header, variable_header, payload)
        instance.protocol_ts = datetime.now()
        return instance

    def __repr__(self):
        return type(self).__name__ + \
                'ts={}, fixed={}, variable={}, payload={}'.format(
                    self.protocol_ts, self.fixed_header,
                    self.variable_header, self.payload
                )
